import os
from tqdm import tqdm
import shelve
"""
    生成反汇编工具类
    所有工具都能进行处理的样本
"""

angrSavePath = r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/angr"
redare2SavePath = r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/radare2_functionSim"
retdecSavePath = r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/retdec"
functionSimPath = r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/functionSim"


class disassemblyTool:
    def __init__(self) -> None:
        pass

    def find_sample_name(self, paths):
        """
            获得路径下都存在的样本名
        """
        size = len(paths)
        datas = []
        for path in paths:
            sampleName = []
            for file in os.listdir(path):
                data = file.split(".")
                if (len(data) == 2 and data[1] == "dir"):
                    sampleName.append(data[0])
            datas.append(sampleName)
            print("{}中处理好的样本数量:\t{}".format(path, len(sampleName)))
        res = []
        for name in datas[0]:
            flag = True
            for i in range(1, size):
                if name not in datas[i]:
                    flag = False
            if flag:
                res.append(name)
        print("{}个路径中都存在的样本:\t{}".format(size, len(res)))
        return res

    def get_used_samples(self):
        """
            获得反汇编实验能够正确预处理的样本
        """
        paths = [angrSavePath, retdecSavePath,
                 redare2SavePath, functionSimPath]
        res = self.find_sample_name(paths)
        return res

    def get_path_sample_inf(self, tarPath):
        """
            获得一个路径下样本的处理信息
            返回：
            {name:节点数}
        """
        g = os.walk(tarPath)
        res = {}
        for path, dir_list, file_list in g:
            for file_name in tqdm(file_list):
                a = file_name.split(".")
                if len(a) == 2 and a[1] == "dir":
                    try:
                        with shelve.open(path+"//"+a[0]) as file:
                            cg = file["cg"]
                            res[a[0]] = len(cg)
                    except Exception as e:
                        print("Exception->{}".format(e))
                        print("path:{}--name:{}".format(tarPath, a[0]))
                        res[a[0]] = 0
        return res

    def get_disassembly_sample(self, minValue=5, maxValue=5000):
        """
            选择节点数量在一定范围内的样本
        """
        useSaveData = True
        paths = [functionSimPath, angrSavePath,
                 retdecSavePath, redare2SavePath]
        sampleName = self.get_used_samples()
        res = {}
        if useSaveData:
            with shelve.open(r"/home/cyw/projects/function_sim_project/disassemblyTools"+"//" + "nodeInf") as file:
                res = file["res"]
        else:
            for path in paths:
                res[path] = self.get_path_sample_inf(path)
            with shelve.open(r"/home/cyw/projects/function_sim_project/disassemblyTools"+"//" + "nodeInf") as file:
                file["res"] = res
                print("样本节点数据保存成功")

        minfailedNum = [0]*len(paths)
        maxfailedNum = [0]*len(paths)

        ans = []
        for name in sampleName:
            flag = 0
            for i in range(len(paths)):
                try:
                    tarNum = res[paths[i]][name]
                    if tarNum < minValue:
                        minfailedNum[i] += 1
                    elif tarNum > maxValue:
                        maxfailedNum[i] += 1
                    else:
                        flag += 1
                except Exception as e:
                    print("保存数据未更新，请重新生成保存数据")
                    print("exception->{}".format(e))
            if (flag == len(paths)):
                ans.append(name)
        for i in range(len(paths)):
            path = paths[i]
            print("{}:\t 低于阈值：{}\t 高于阈值：{}\n".format(
                path, minfailedNum[i], maxfailedNum[i]))
        print("最终的样本数量为：{}".format(len(ans)))
        return ans

    def get_path_inf(self, path, savename):
        """
            获得如下信息：
                平均节点数
                平均边数
                各类函数数量
        """
        names = []
        nodes = []
        edges = []
        functypes = []

        g = os.walk(path)
        for path, dir_list, file_list in g:
            for file_name in tqdm(file_list):
                a = file_name.split(".")
                if len(a) == 2 and a[1] == "dir":
                    try:
                        with shelve.open(path+"//"+a[0]) as file:
                            if savename != "ida":
                                functype = file["functype"]
                            else:
                                functype = file["funcType"]
                            cg = file["cg"]
                            names.append(a[0])
                            size = len(cg)
                            nodes.append(size)
                            num = 0
                            for i in range(size):
                                for j in range(size):
                                    if cg[i][j] != 0:
                                        num += 1
                            edges.append(num)

                            a = [0, 0, 0]
                            for ind in functype:
                                if (functype[ind] == "local"):
                                    a[0] += 1
                                elif functype[ind] == "dynamic import":
                                    a[1] += 1
                                else:
                                    a[2] += 1
                            functypes.append(a)
                    except Exception as e:
                        print("Exception->{}".format(e))
                        print("path:{}--name:{}".format(path, a[0]))
        size = len(names)
        res = {}
        res["names"] = names
        res["nodes"] = nodes
        res["edges"] = edges
        res["functypes"] = functypes
        sum = 0
        for node in nodes:
            sum += node
        res["avgnodes"] = sum/size
        sum = 0
        for edge in edges:
            sum += edge
        res["avgedge"] = sum/size
        a = [0, 0, 0]
        for i in range(size):
            for j in range(3):
                a[j] += functypes[i][j]

        res["avgfunctype"] = [a[0]/size, a[1]/size, a[2]/size]
        with shelve.open(r"/home/cyw/projects/function_sim_project/disassemblyTools/tempData"+"//" + savename) as file:
            file["res"] = res
            print("数据保存成功")

    def gene_sample_infs(self):
        """
            获得处理后样本的基础信息
        """
        print("正在获得各个反汇编工具处理后的样本基础信息")
        print("之前报错的代码错误")
        paths = [functionSimPath, angrSavePath,
                 retdecSavePath, redare2SavePath]
        # saveName=["angr","retdec","radare2","ida"]
        # 这里标签对应错了。。。。
        saveName = ["ida", "angr", "retdec", "radare2"]
        for i in range(4):
            self.get_path_inf(paths[i], saveName[i])

    def get_gene_sample_infs(self, names):
        """
            获得names里面的样本统计信息
        """
        for saveName in ["angr", "retdec", "radare2", "ida"]:
            with shelve.open(r"/home/cyw/projects/function_sim_project/disassemblyTools/tempData"+"//" + saveName) as file:
                res = file["res"]
                nodes = res["nodes"]
                edges = res["edges"]
                functypes = res["functypes"]
                tarNames = res["names"]
                NodeNum, edgeNum = 0, 0
                a, b, c = 0, 0, 0
                for ind in range(len(tarNames)):
                    if tarNames[ind] in names:
                        NodeNum += nodes[ind]
                        edgeNum += edges[ind]
                        a += functypes[ind][0]
                        b += functypes[ind][1]
                        c += functypes[ind][2]
                size = len(names)
                print("{}:样本数量：{}".format(saveName, size))
                print("平均节点数：{:.2f}\t平均边数：{:.2f}\t本地函数:{:.2f}\t动态导入:{:.2f}\t静态链接:{:.2f}"
                      .format(NodeNum/size, edgeNum/size, a/size, b/size, c/size))


if __name__ == "__main__":
    tools = disassemblyTool()
    names = tools.get_disassembly_sample(5, 5000)
    tools.gene_sample_infs()
    tools.get_gene_sample_infs(names)
